import sys
import pathlib

def readDB(link):
    _d = {}
    with open(link) as infile:
        for i in infile:
            el = i.strip().split('\t')
            if len(el) == 2:
                _d[el[0]] = el[1]
            else:
                _d[el[0]] = ''
    return _d

atb = readDB('/data/reference/pathoseq/card_2021/bin/antibiotic.txt')

atbr = readDB('/data/reference/pathoseq/card_2021/bin/antibiotic-resistant.txt')

mn = readDB('/data/reference/pathoseq/card_2021/bin/microbio-name.txt')

n = {}
raw_n = {}
ns = {}
p = pathlib.Path(sys.argv[1])
name = p.name.replace('.gene_mapping_data.raw', '')
with open(sys.argv[1]) as rgi:
    rgi.readline()
    for i in rgi:
        tmp = i.strip().split('\t')
        ns[tmp[0]] = float(tmp[9])

        antibiotic = '-'
        if tmp[24].count('; ') > 0:
            line2 = tmp[24].split('; ')
            for key, value in enumerate(line2):
                if atb.get(value):
                    line2[key] = atb[value]
            antibiotic = '; '.join(line2)
        else:
            if atb.get(tmp[24]):
                antibiotic = atb[tmp[24]]

        atb_resistant = '-'
        if atbr.get(tmp[25]):
            atb_resistant = atbr[tmp[25]]
        elif tmp[25].count('; '):
            resis = tmp[25].split('; ')[0]
            if atbr.get(resis):
                atb_resistant = atbr[resis]
        
        microbio = tmp[8]
        if mn.get(tmp[8]):
            microbio = mn[tmp[8]]+tmp[8]
        
        n[tmp[0]] = '\t'.join((name,tmp[0],tmp[9],tmp[12],antibiotic,atb_resistant,microbio,tmp[24],tmp[25],tmp[8]))
        raw_n[tmp[0]] = '\t'.join((name,tmp[0],tmp[9],tmp[12],tmp[24],tmp[25],tmp[8]))
    
    nss = sorted(ns.items(), key=lambda x: x[1], reverse=True)
    with open(p.parent / pathlib.Path(f'{name}.gene_mapping_data.xls'), 'w', encoding='gbk') as out:
        out.write("ID\tgene\tcount\tcoverage\tdrug_cn\tmed_cn\tspecies_cn\tdrug_en\tmed_en\tspecies_en\n")
        for i in nss:
            out.write(n[i[0]]+'\n')
    with open(p.parent / pathlib.Path(f'{name}.gene_mapping_data.txt'), 'w') as out2:
        out2.write("ID\tgene\tcount\tcoverage\tdrug\tmed\tspecies\n")
        for i in nss:
            cover = float(raw_n[i[0]].split('\t')[3])
            if cover > 50:
                out2.write(raw_n[i[0]]+'\n')
